import rateLimit from 'koa-ratelimit';
import { Context } from 'koa';
import { config } from '@/config';
import { getClientIP } from '@/utils';

/**
 * 内存存储（简单实现）
 */
const db = new Map();

/**
 * 基础限流中间件
 */
export const baseLimiter = rateLimit({
  driver: 'memory',
  db,
  duration: config.rateLimit.windowMs, // 时间窗口（毫秒）
  errorMessage: '请求过于频繁，请稍后再试',
  max: config.rateLimit.maxRequests, // 最大请求数
  disableHeader: false,
  id: (ctx: Context) => getClientIP(ctx.request),
  headers: {
    remaining: 'Rate-Limit-Remaining',
    reset: 'Rate-Limit-Reset',
    total: 'Rate-Limit-Total',
  },
} as any);

/**
 * 严格限流中间件（用于敏感操作）
 */
export const strictLimiter = rateLimit({
  driver: 'memory',
  db,
  duration: 15 * 60 * 1000, // 15分钟
  errorMessage: '操作过于频繁，请稍后再试',
  max: 5, // 15分钟内最多5次
  disableHeader: false,
  id: (ctx: Context) => getClientIP(ctx.request),
  headers: {
    remaining: 'Rate-Limit-Remaining',
    reset: 'Rate-Limit-Reset',
    total: 'Rate-Limit-Total',
  },
} as any);

/**
 * 登录限流中间件
 */
export const loginLimiter = rateLimit({
  driver: 'memory',
  db,
  duration: 15 * 60 * 1000, // 15分钟
  errorMessage: '登录尝试过于频繁，请稍后再试',
  max: 5, // 15分钟内最多5次登录尝试
  disableHeader: false,
  id: (ctx: Context) => {
    // 使用 IP + 用户名的组合作为标识
    const ip = getClientIP(ctx.request);
    const email = ctx.request.body?.email || '';
    return `${ip}:${email}`;
  },
  headers: {
    remaining: 'Rate-Limit-Remaining',
    reset: 'Rate-Limit-Reset',
    total: 'Rate-Limit-Total',
  },
});

/**
 * API 限流中间件
 */
export const apiLimiter = rateLimit({
  driver: 'memory',
  db,
  duration: 60 * 1000, // 1分钟
  errorMessage: 'API调用过于频繁，请稍后再试',
  max: 60, // 1分钟内最多60次
  disableHeader: false,
  id: (ctx: Context) => {
    // 认证用户使用用户ID，未认证用户使用IP
    if (ctx.state.user?.id) {
      return `user:${ctx.state.user.id}`;
    }
    return `ip:${getClientIP(ctx.request)}`;
  },
  headers: {
    remaining: 'Rate-Limit-Remaining',
    reset: 'Rate-Limit-Reset',
    total: 'Rate-Limit-Total',
  },
});

/**
 * 创建自定义限流中间件
 */
export const createLimiter = (options: {
  duration: number;
  max: number;
  message?: string;
  keyGenerator?: (ctx: Context) => string;
}) => {
  return rateLimit({
    driver: 'memory',
    db,
    duration: options.duration,
    errorMessage: options.message || '请求过于频繁，请稍后再试',
    max: options.max,
    disableHeader: false,
    id: options.keyGenerator || ((ctx: Context) => getClientIP(ctx.request)),
    headers: {
      remaining: 'Rate-Limit-Remaining',
      reset: 'Rate-Limit-Reset',
      total: 'Rate-Limit-Total',
    },
  });
};

export default baseLimiter;